% findFringeParams.m
%
% author: Kevin O'Holleran ko311@cam.ac.uk
%
% Description: A function that accepts two complex matrices and finds the
% optimal correlation vector between the two matrices across a given search
% width. It does this by applying cosines to the Fourier transform of one
% matrix (shiting the matrix laterally in it's original space). This is
% done to achieve sub-pixel image registration.
%
% The optimisation step uses a discrete correlation maximum position in
% order to start sub-pixel searching in the neighbourhood of the real
% solution. In this neighbourhood a 10x10 grid is searched and then the
% maximum is located and the grid is zoomed onto that location by a factor
% of 2. This continues for n iterations (resulting in 2^n zooms, each time
% the sub-pixel accuracy becoming better by a factor of 2).

function [r_est a_est] = findFringeParams(k1,kp1,search_width,otf0)

[height width] = size(k1);

if nargin <3
    search_width = width/4;
    otf0 = ones(height,width);
elseif nargin <4
    otf0 = ones(height,width);
end

% discrete image correlation
xs = repmat(1:(width),(width),1);
ys = repmat(transpose(1:(width)),1,(width));
d1 = sample(k1,2);
d2 = sample(kp1,2);
cc = corr2dscan(abs(fftshift(d1)),abs(fftshift(d2)),search_width);
% central region masked with circular block to remove possible high
% correlations near the origin.
cc = mask(cc,0.7,true,false);
[m mi] = max(cc(:));
mx = floor((mi-1)/search_width)+1;
my =  rem((mi-1),search_width)+1;

% start searching using Fourier technique to achieve sub-pixel image
% registration
nt = 4;
rg = 2*sqrt((mx-search_width/2)^2 + (my-search_width/2)^2);
ag = atan2((my-search_width/2),(mx-search_width/2));
psf = ifft2(fftshift(otf0));
%% For checking correlations
for t = 0:nt
    display(['Fringe registration iteration ' int2str(t+1) ' of ' int2str(nt+1)])
    an = 0.04/(2^t);
    rn = 4/(2^t);
    n = 10;
    a = (ag-an/2):an/n:(ag+an/2);
    r = (rg-rn/2):rn/n:(rg+rn/2);
    cc = zeros(numel(a),numel(r));
    box = 6;
    for i = 1:numel(a)
        ai = -pi+a(i);
        for j = 1:numel(r)
            w = r(j);
            kp1r = ifft2(kp1);
            shiftp1 = exp(1i*2*pi*w*(cos(ai).*xs+sin(ai).*ys)/width);
            kpt1 = fft2(kp1r.*shiftp1);
            otf0_shifted = fft2(psf.*shiftp1);
            A = fftshift(kpt1);
            %A = A((width/2-box):(width/2+box),(height/2-box):(height/2+box));
            A = A.*otf0;
            B = fftshift(k1);
            B = B.*otf0_shifted;
            %B = B((width/2-box):(width/2+box),(height/2-box):(height/2+box));
            cc(i,j) = corr2d(abs(A),abs(B));
        end
    end
    [m mi] = max(cc(:));
    mx = floor((mi-1)/numel(a))+1;
    my =  rem((mi-1),numel(a))+1;
    r_est = r(mx);
    a_est = a(my);
    rg = r_est;
    ag = a_est;
end
display(['Fringe spatial frequency: ' num2str(r_est)])
display(['Fringe angle: ' num2str(a_est)])
end

